#include "CoordinateProjectionUtils.h"
#include "List.h"
#include "Constraint.h"
#include "Utils.h"
#include "LU.h"
#include "Svd.h"
#include "MultibodySystem.h"

using namespace RSIM;

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::resize_map(const int& n, 
                                         List <List <int> >& bmap)const{
        for(int i=0;i<bmap.getLength();++i)
                bmap[i].clear();
        bmap.clear();
        List <int> tmp;
        for(int i=0;i<n;++i)
                bmap.addNode(tmp);
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::sortSysId(const List <List<int> >& ccouple,
                                          const List <Constraint*>& constraints,
                                          List <List <int> >& sid)const{
                                                  
        // This differs form the actual number of constraints.
        // nConstraints is number of rows of ccouple_ or includes
        // the information for the constraints which must be solved
        // simultaneously.                                          
        const int& nConstraints = ccouple.getLength();        
        if(nConstraints){                
                for(int i=0;i<nConstraints;++i){
                        const List <int>& coup_cons = ccouple[i];
                        List <int> sidList_tmp;
                        for(int j=0;j<coup_cons.getLength();++j){
                                
                                const Constraint *ptr = constraints[coup_cons[j]];
                                
                                // Get Parent and Child SysId
                                const int& pid = ptr->getParentSysId();
                                const int& cid = ptr->getChildSysId();
                                
                                // Check if parent Id is already added
                                int tmp = sidList_tmp.findElement(pid);
                                if(!(tmp+1)){
                                        sidList_tmp.addNode(pid);
                                }
                                
                                // Check if child Id is already added
                                tmp = sidList_tmp.findElement(cid);
                                if(!(tmp+1)){
                                        sidList_tmp.addNode(cid);
                                }
                        }
                        sidList_tmp.finalize();
                        sid.addNode(sidList_tmp);
                        sidList_tmp.clear();
                }
                sid.finalize();  
        }//if(nConstraints)        
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::calcPositionLevelErrorForAllConstraints(List <Constraint*>& constraints)const{
        const int& nConstraints = constraints.getLength();
        
        for(int i=0;i<nConstraints;++i){
                // i^th constraint
                Constraint *cPtr = constraints[i];
                
                // Check if the constraint is imposed on linear velocity
                const bool lin_const = cPtr->imposeConstraintOnLinearVelocity();
                
                // If linear velocity is constrained, then one needs to 
                // calculate the cumulative betak                                
                if(lin_const)                                      
                        cPtr->calcCumulativeBeta();
                
                // Now that beta has been updated, position level errors can be computed.
                cPtr->calcPositionLevelError();
        }
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::popuplateCqVector(      const List<Constraint*>& constraint, 
                                                        const List <List<int> >& ccouple,
                                                        Vector *Cq){
        const int& ncc = ccouple.getLength();
        for(int i=0;i<ncc;++i){ 
                double *cq_ptr = Cq[i].wPtr();
                const List <int>& coup_cons = ccouple[i];
                
                int m = 0;
                for(int j=0;j<coup_cons.getLength();++j){
                        const Constraint *cPtr = constraint[coup_cons[j]];
                        const int nalg = cPtr->getNAlgebraicConstraints();
                        
                        const double *pos_error = cPtr->getPositionErrorAsVector().Ptr();
                        
                        for(int k=0;k<nalg;++k){
                                cq_ptr[m++] = pos_error[k];
                        }

                }//for(int j=0;j<coup_cons.getLength()
        }//for(int i=0;i<ncc
        
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::calcPositionErrorNorms(const Vector *Cq, 
                                                       const int& nccouple, 
                                                       Vector& PositionErrorNorms)const{
        for(int i=0;i<nccouple;++i)                
                PositionErrorNorms[i] = Utils::norm(Cq[i]);
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::calcDelQ(const Matrix *CJac,
                                         const Vector *CQ,
                                         Vector *delq)const{
        
        // This solves the constrained minimization problem
        // minimize (delq^T * delq) subject to constraint 
        // Jac * delq = -CQ 
                                                 
        // New labels
        const Matrix& Jac = *CJac;
        const Vector& Cq = *CQ;
                
        Vector& dq = *delq;
        
        // The constraint Jacobian could be rank deficient but the 
        // singular value decomposition will always exist.
        SVD svd(Jac);
        Matrix M;
        
        // Calculate the column basis of the the constraint Jacobian.
        svd.getColumnBasis(M);
        const int& ncol = M.getNCol();
        
        // Project Cq onto the range of the constraint Jacobian.
        // At this point the equality constraint has been modified
        // to Jac * delq = -M*(~M)*Cq. Note that in case the the constraint
        // jacobian is full rank, M*(~M) is an identity matrix and this calculation
        // is excess. 
        Vector mCq = Vector(M*(~M)*Cq);
                
        Matrix U,S,V;        
        
        // Here Jac = U*S*(~V)
        svd.getU(U);
        svd.getS(S);
        svd.getV(V);
        
        // Equality constraint now is S*(~V)*delq = -(~U)*M*(~M)*Cq
        // Solution to this always exists since we have already ensured that
        // the rhs is in the range of lhs
        Vector rhs_tmp = Vector((~U)*mCq);
        Matrix SVT = S*(~V);        
        
        // Ignore the equations which correspond to the zero singular values.
        Vector rhs = rhs_tmp.extractNElements(ncol);        
        Matrix RedJac = SVT.extractNRows(ncol);                        
        
        // The equality constraint at this point is 
        // RedJac * delq = reduced(-(~U)*M*(~M)*Cq)
        Matrix A = RedJac*(~RedJac);        
        LU lu(A);                                           
        Vector B = lu.solve(rhs);
        dq = -((~RedJac)*B);        
}

/////////////////////////////////////////////////////// 

void CoordinateProjectionUtils::updCoupledConstraintQs(const Vector& delq, 
                                                       const List<int>& coup_cons,
                                                       List <Constraint*>& constraint)const{
        // This is the q being updated
        int qid = 0;
        
        const int& nconstraints = coup_cons.getLength();
        for(int i=0;i<nconstraints;++i){
                Constraint *cptr = constraint[coup_cons[i]];
                cptr->updQ(delq, qid);
        }
}

/////////////////////////////////////////////////////// 